import time
from functools import wraps

from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser

from langchain_workflow.llm import ollama_llm
from langchain_workflow.models import (
    ProbabilityResultModel,
    ProductModel,
    ProductRemovePromotionalContentModel,
)
from langchain_workflow.retriever import compression_retriever, rerank
from logger import logger


def timing_decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        logger.info(f"函数 {func.__name__} 执行时间: {execution_time:.4f} 秒")
        return result

    return wrapper


@timing_decorator
def rag(question: str):
    answer = rerank(question, compression_retriever)
    return answer


@timing_decorator
def remove_promotional_content(
    product_model: ProductModel,
) -> ProductRemovePromotionalContentModel:
    prompt_template = """
        -Goal-
        You are an AI assistant tasked with removing promotional and subjective content from product descriptions. 
        Your goal is to retain only objective, factual information about the product. Please follow these guidelines:

        -Notice-
        1. Remove any subjective claims about quality, effectiveness, or superiority.
        2. Eliminate promotional language, marketing slogans, and sales pitches.
        3. Delete statements about safety margins or dependability unless they are quantifiable.
        4. If the product title and description mention the target audience, please retain that information.
        5. Please provide only the rewritten text, without any unnecessary words.
        
        -- Format Instructions --
        {format_instructions}

        -- Input --
        {ori_text}
        """

    parser = JsonOutputParser(pydantic_object=ProductRemovePromotionalContentModel)

    prompt = PromptTemplate(
        template=prompt_template,
        input_variables=["ori_text"],
        partial_variables={"format_instructions": parser.get_format_instructions()},
    )

    chain = prompt | ollama_llm | parser
    max_retries = 3
    retry_count = 0
    while retry_count < max_retries:
        try:
            result = chain.invoke({"ori_text": product_model.ori_text})
            res = ProductRemovePromotionalContentModel(**result)
            break
        except Exception as e:
            logger.warn(f"尝试 {retry_count + 1} 失败: {e}")
            retry_count += 1
            if retry_count == max_retries:
                logger.error(f"达到最大重试次数 {max_retries}，返回默认结果")
                res = ProductRemovePromotionalContentModel(
                    pd_title=product_model.pd_title, pd_desc=product_model.pd_desc
                )
    return res


@timing_decorator
def get_product_status(
    product_model: ProductRemovePromotionalContentModel,
    prohibited_policy_str: str,
) -> ProbabilityResultModel:
    parser = JsonOutputParser(pydantic_object=ProbabilityResultModel)

    with open("langchain_workflow/prompt.md") as f:
        prompt_template = f.read()

    prompt = PromptTemplate(
        template=prompt_template,
        input_variables=["product_info", "polices"],
        partial_variables={"format_instructions": parser.get_format_instructions()},
    )
    chain = prompt | ollama_llm | parser
    max_retries = 3
    retry_count = 0
    #如果超过重试次数3次，禁止标签和合规标签默认返回0，判断理由均规定为处理失败
    result = ProbabilityResultModel(
        is_prohibited_probability=0,
        requires_compliance_check_probability=0,
        is_prohibited_reason="处理失败",
        requires_compliance_check_reason="处理失败",
    )
    #设置最大重试次数，最大次数为3次
    while retry_count < max_retries:
        try:
            result_dict = chain.invoke(
                {
                    "product_info": product_model.rag_text,
                    "polices": prohibited_policy_str,
                }
            )
            #将字典转换为ProbabilityResultModel对象，对返回结果进行处理
            return ProbabilityResultModel(**result_dict)
        except Exception as e:
            logger.warn(f"尝试 {retry_count + 1} 失败: {e}")
            retry_count += 1
            if retry_count == max_retries:
                logger.error(f"达到最大重试次数 {max_retries}，返回默认结果")
                return result

    return result
